# 导入必要的包
# 导入必要的包
import pandas as pd
import statsmodels.formula.api as smf
from sqlalchemy import create_engine
import pymysql
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.decomposition import PCA

# 数据库配置
db_config = {
    'host': 'localhost',
    'user': 'wsj',
    'password': '123456',
    'database': 'tushare',
    'port': 3307,
    'charset': 'utf8mb4'
}

# 创建数据库连接引擎
engine = create_engine(
    f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}?charset={db_config['charset']}"
)

# 创建 PyMySQL 连接
conn = pymysql.connect(**db_config)

# 检查数据库连接
if conn.open:
    print("数据库连接成功")
else:
    print("数据库连接失败")
    exit()

# 设置分批读取的大小
chunk_size = 10000

# 获取华夏银行日线数据
query = """
SELECT d.*, i.closes as i_closes, i.vol as i_vol
FROM date_1 d
LEFT JOIN index_daily i ON d.trade_date = i.trade_date AND i.ts_code = '000001.SH'
WHERE d.trade_date BETWEEN '2023-01-01' AND '2023-12-31' AND d.ts_code = '600015.SH'
"""

# 读取数据
chunks = pd.read_sql_query(query, conn, chunksize=chunk_size)
df = pd.concat(chunks, ignore_index=True)

# 检查数据框是否为空
if df.empty:
    print("数据框为空，无法继续")
    exit()
else:
    print("数据框形状：", df.shape)
    print("数据框前几行：", df.head())

# 计算收益率和成交量变化率
df['zd_closes'] = df['closes'].pct_change()
df['zs_closes'] = df['i_closes'].pct_change()
df['zs_vol'] = df['i_vol'].pct_change()

# 去除缺失值
df = df.dropna(subset=['zd_closes', 'zs_closes', 'zs_vol'])

# 检查数据框是否为空
if df.empty:
    print("去除缺失值后数据框为空，无法继续")
    exit()
else:
    print("去除缺失值后数据框形状：", df.shape)

# 特征选择
features = ['vol', 'amount', 'buy_sm_vol', 'buy_sm_amount', 'sell_sm_vol', 'sell_sm_amount', 
            'buy_md_vol', 'buy_md_amount', 'sell_md_vol', 'sell_md_amount', 
            'buy_lg_vol', 'buy_lg_amount', 'sell_lg_vol', 'sell_lg_amount', 
            'buy_elg_vol', 'buy_elg_amount', 'sell_elg_vol', 'sell_elg_amount', 'net_mf_vol']

# 检查特征列是否存在
print("特征列：", features)
print("数据框列：", df.columns.tolist())

# 检查数据类型
print("数据框数据类型：", df.dtypes)

# 一元线性回归
formula_1 = 'zd_closes ~ zs_closes'
if not df.empty:
    model_1 = smf.ols(formula_1, data=df).fit()
    print("一元线性回归结果：")
    print(model_1.summary())

    # 绘制一元线性回归散点图
    plt.figure(figsize=(10, 6))
    plt.scatter(model_1.fittedvalues, df['zd_closes'], alpha=0.7)
    plt.title('一元线性回归拟合值 vs 实际值')
    plt.xlabel('拟合值')
    plt.ylabel('实际值')
    plt.grid(True)
    plt.show()
else:
    print("数据框为空，无法构建一元线性回归模型")



# 多元线性回归
formula_2 = 'zd_closes ~ ' + ' + '.join(features)
if not df.empty:
    model_2 = smf.ols(formula_2, data=df).fit()
    print("\n多元线性回归结果：")
    print(model_2.summary())

    # 绘制多元线性回归散点图
    plt.figure(figsize=(10, 6))
    plt.scatter(model_2.fittedvalues, df['zd_closes'], alpha=0.7)
    plt.title('多元线性回归拟合值 vs 实际值')
    plt.xlabel('拟合值')
    plt.ylabel('实际值')
    plt.grid(True)
    plt.show()
else:
    print("数据框为空，无法构建多元线性回归模型")

# 主成分分析
if not df.empty:
    pca = PCA(n_components=0.95)  # 保留95%的方差
    pca_features = pca.fit_transform(df[features])
    explained_variance = pca.explained_variance_ratio_
    print("\n主成分分析结果：")
    print(f"主成分个数：{pca.n_components_}")
    print(f"各主成分解释的方差比例：{explained_variance}")

    # 绘制主成分分析累计解释方差比例图
    plt.figure(figsize=(10, 6))
    plt.plot(np.cumsum(explained_variance), marker='o')
    plt.xlabel('主成分个数')
    plt.ylabel('累计解释方差比例')
    plt.title('主成分分析累计解释方差比例')
    plt.grid(True)
    plt.show()
else:
    print("数据框为空，无法进行主成分分析")

# 绘制特征相关性热力图
if not df.empty:
    plt.figure(figsize=(12, 10))
    correlation_matrix = df[features].corr()
    sns.heatmap(correlation_matrix, annot=True, cmap='coolwarm', cbar=True, fmt='.2f', linewidths=.5)
    plt.title('特征相关性热力图')
    plt.show()
else:
    print("数据框为空，无法绘制热力图")